[muse_glimmer] Enable CP for the multimodal model - #4387
Conversation
meta-cla
Bot
added
the
CLA Signed
Build packed vision source indices before context-parallel token sharding and carry them through the same load-balancing permutation. Keep vision compute replicated across CP, then use a stateless sharding boundary to gather each local token shard while preserving TP and FSDP gradient semantics. Generalize the shared vision attention layouts, keep multimodal embedding preparation on pipeline stage 0, and support the combined TP+CP+PP+SP topology without changing existing non-CP behavior or checkpoint keys. Test Plan: 71 focused CPU tests; 2 distributed GPU tests; 10-step CP2xPP2 and TP2xCP2xPP2+SP real-process-group runs; exact 10-step deterministic loss and grad-norm parity; repository pre-commit hooks excluding the separately documented optional-dependency Pyrefly blocker.
| return vision_bank_indices_T.masked_fill(~vision_mask_T, -1) | ||
|
|
||
|
|
||
| class VisionScatter(Module): |
There was a problem hiding this comment.
Question 1: are you putting it here because other multimodal models could share? Is there a vlm model in torchtitan that cannot use this to implement CP?
Question 2: do you create this stateless module only because you want _vision_scatter_config to perform spmd collectives at the module boundary? Fwiw we are moving away from such pattern because they couldn't express fused comm + computation. Could you put collectives INSIDE model code instead? I think we at least we could have a function that does vision scattering, instead of a module.
There was a problem hiding this comment.
- In principle, this should allow CP on any VLM that has support for CP in the decoder.
- I will try to remove that.
There was a problem hiding this comment.
@tianyu-l We are moving away from doing spmd collective at the module boundary. But we currently can only swap modules. When you said, calling collectives inside model code, is this still a swappable one or a if/else condition inside the model code?
There was a problem hiding this comment.
There is only one redistribute required for this, where we SP shard vision_bank_indices_T. So this can be done in preprocess_inputs which already has parallel_dims and parallelism. So the scatter function does not need to do redistribute (note that SP sharding is also compatible with scattering using vision_bank_indices_T).
But I am not sure how this kind of parallel_dims dependent redistribute can be done without a module. Modules can remember it from parallelize like how it is done in ScaledBiasRowwiseLinear.
There was a problem hiding this comment.
But we currently can only swap modules.
Didn't understand
When you said, calling collectives inside model code, is this still a swappable one or a if/else condition inside the model code?
Didn't understand either, lol. I think the latest commit LGTM.
But I am not sure how this kind of parallel_dims dependent redistribute can be done without a module. Modules can remember it from parallelize like how it is done in ScaledBiasRowwiseLinear.
That's good question. I think we should
- either have spmd context provide the parallel info
- or passing parallel info along the the model definition
I would prefer 1, if it can be done properly (thread local is good enough??) cc @pianpwk
76d3e8d to
ee49a57
Compare
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
do you plan to support other VLM soon?
| ) | ||
| if parallelism.spmd_backend == "spmd_types": | ||
| if ( | ||
| parallelism.enable_sequence_parallel |
|
can qwen3.5 vl also benefit from the common utils here? When looking into qwen3.5 CP, I found that we need to first do vision scatter -> CP shard. Is this what you are doing here? |
4 participants
Summary
This PR enables CP for the multimodal muse_glimmer.
MuseGlimmerModel.preprocess_inputsbuildsvision_bank_indices_T: [T]from unsharded text tokens, which maps each token position to the index of vision encoder output that should be placed there.Then,
vision_bank_indices_Tgets permuted and sharded byprepare_context_parallel_inputjointly with tokens, so it can be used in the model forward to gather the vision encoder output (which is replicated on CP) after token is embedded.Test Plan
TP2+SP vs TP2+CP2+SP: max loss difference 1.38e-5; max grad_norm difference 5.04e-4 from running
loss_compare.py